package com.tfswx.ocrengine.util;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;

/**
 * @author Ricky
 * @date 2021/5/14
 */
@Slf4j
public class HeheOcrUtils {
    private HeheOcrUtils() {
    }

    public static JSONObject convert2Ali(JSONArray documents) {
        JSONObject aliDataSource = new JSONObject();
        aliDataSource.put("prism_version", "");
        JSONObject pageJsonObject = documents.getJSONObject(0);
        aliDataSource.put("angle", pageJsonObject.getDoubleValue("rotate_angle"));
        aliDataSource.put("height", pageJsonObject.getIntValue("height"));
        aliDataSource.put("orgHeight", pageJsonObject.getIntValue("orgHeight"));
        aliDataSource.put("orgWidth", pageJsonObject.getIntValue("orgWidth"));
        aliDataSource.put("width", pageJsonObject.getIntValue("width"));
        JSONArray prismWordsInfo = parseBlocks(pageJsonObject.getJSONArray("blocks"));

        aliDataSource.put("prism_wnum", prismWordsInfo.size());
        StringBuilder content = new StringBuilder();
        for (int i = 0; i < prismWordsInfo.size(); i++) {
            String word = prismWordsInfo.getJSONObject(i).getString("word");
            if (i != 0) {
                content.append("\n");
            }
            content.append(word);
        }
        aliDataSource.put("content", content.toString());
        aliDataSource.put("prism_wordsInfo", prismWordsInfo);
        aliDataSource.put("prism_tablesInfo", new JSONArray());
        return aliDataSource;
    }

    private static JSONArray parseBlocks(JSONArray blocks) {
        JSONArray prismWordsInfo = new JSONArray();

        JSONObject prismWord;
        for (int i = 0; i < blocks.size(); i++) {
            prismWord = new JSONObject();
            prismWord.put("paragraphId", 0);
            prismWord.put("rowId", 0);
            JSONObject block = blocks.getJSONObject(i);
            // 获取位置
            JSONArray pos = block.getJSONObject("position").getJSONArray("vertices");
            prismWord.put("pos", pos);

            JSONObject line = block.getJSONArray("lines").getJSONObject(0);

            // 获取分数
            double score = line.getDoubleValue("score") * 100;
            prismWord.put("prob", score);
            prismWord.put("word", line.getString("text"));

            // 获取charInfo
            JSONArray charInfos = new JSONArray();
            JSONArray characters = line.getJSONArray("characters");
            JSONObject charInfo;
            for (int i1 = 0; i1 < characters.size(); i1++) {
                charInfo = new JSONObject();
                JSONObject character = characters.getJSONObject(i1);
                charInfo.put("prob", character.getDoubleValue("score") * 100);
                charInfo.put("word", character.getString("text"));
                JSONObject bound = character.getJSONObject("position").getJSONObject("bounding_box");
                charInfo.put("w", bound.getIntValue("width"));
                charInfo.put("x", bound.getIntValue("left"));
                charInfo.put("h", bound.getIntValue("height"));
                charInfo.put("y", bound.getIntValue("top"));

                charInfos.add(charInfo);
            }
            prismWord.put("charInfo", charInfos);

            prismWordsInfo.add(prismWord);
        }

        return prismWordsInfo;
    }
}
